import argparse
import osutil


def main():
    parser = argparse.ArgumentParser(prog='folder2tree', description="将指定目录以字符树的结构呈现")
    parser.add_argument("root", type=str, help="根目录路径")
    parser.add_argument("-l", '--limit', type=int, default=1000, help="输出最大行数，若字符树所占行数超出该值，则按该值截断")
    parser.add_argument("-e", "--exclude", type=str, default=None, help="指定pattern用于排除文件，多个表达式用逗号分隔。支持类似正则表达式的模式串，不过 `.` 被视为普通字符，`?` 匹配单个字符，`*` 匹配多个字符")
    parser.add_argument("-d", "--directory-only", action="store_true", help="只显示目录")
    args = parser.parse_args()
    root = args.root
    length_limit = args.limit
    if args.exclude:
        exclude_repath = args.exclude.replace(',', '$|').replace('.', r'\.').replace('?', '.').replace('*', '.*?') + '$'
    else:
        exclude_repath = None
    
    try:
        num_dirs, num_files, result, is_over = osutil.tree(root, exclude_repath=exclude_repath, length_limit=length_limit, only_directories=args.directory_only)
    except Exception as e:
        print(e)
        exit(-1)
    if result:
        result += '\n'
    if not is_over:
        result += f'... 结果超过 {length_limit} 行，已截断\n'
    result += f'已显示 {num_dirs} 个文件夹' + f', {num_files} 个文件'
    print(result)

if __name__ == "__main__":
    main()
